﻿using System;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;

namespace Mars.GpuInstance
{
    public abstract class VariationBufferData
    {
        public string name { get; set; }
        public abstract void SetBuffer(ref ComputeBuffer buffer, NativeList<int> activeIdList);
        public abstract void Clear();
    }

    public class VariationBufferData<T> : VariationBufferData where T : unmanaged
    {
        readonly Dictionary<int, T> idToDataMap = new();

        public int Count => idToDataMap.Count;

        public void AddData(int id, T val)
        {
            idToDataMap[id] = val;
        }

        public void RemoveData(int id)
        {
            idToDataMap.Remove(id);
        }

        public override void SetBuffer(ref ComputeBuffer buffer, NativeList<int> activeIdList)
        {
            var len = activeIdList.Length;
            var tempArray = new NativeArray<T>(len, Allocator.Temp);
            for (var i = 0; i < len; ++i)
            {
                var id = activeIdList[i];
                if (idToDataMap.TryGetValue(id, out var v))
                {
                    tempArray[i] = v;
                }
            }

            var stride = UnsafeUtility.SizeOf<T>();
            if (buffer != null && buffer.IsValid() && buffer.count >= len && buffer.stride == stride)
            {
                buffer.SetData(tempArray, 0, 0, len);
            }
            else
            {
                if (buffer != null)
                {
                    buffer.Dispose();
                }

                buffer = new ComputeBuffer(len > 8 ? len : 8, stride);
                buffer.SetData(tempArray, 0, 0, len);
            }
            tempArray.Dispose();
        }

        public override void Clear()
        {
            idToDataMap.Clear();
        }
    }
}
